#!/bin/bash

# --- 默认配置 ---
MYSQL_USER="root"
MYSQL_PWD="test@2025"
SYSBENCH_EXECUTABLE="/usr/local/bin/sysbench"
LUA_SCRIPT_BASE_PATH="/root/sysbench_auto_opt/sysbench-1.0.20/src/lua" 
DB_NAME_PREFIX="sysbench"
SYSBENCH_OUTPUT_DIR="./sysbench_results"

# --- 函数：显示用法 ---
usage() {
    echo "用法: $0 -p <端口列表> -S <场景列表> [--prepare] [sysbench参数...]"
    echo ""
    echo "脚本控制参数:"
    echo "  -p <列表>      用逗号分隔的 MySQL 实例端口列表 (必需)。"
    echo "  -S <列表>      用逗号分隔的 Sysbench 场景列表 (必需)。"
    echo "  --prepare      (标志) 如果指定，则运行 sysbench prepare 阶段。"
    echo ""
    echo "Sysbench 参数:"
    echo "  所有其他以 '--' 开头的参数都将直接传递给 sysbench 命令。"
    echo "  例如: --tables=32 --table-size=5000000 --threads=256 --time=120 --report-interval=10"
    echo "  如果未提供，sysbench 将使用其自身的默认值（例如 --time=10）。"
    echo ""
    echo "示例:"
    echo "  # 准备数据 (32个表, 每个500万行), 使用64个线程"
    echo "  $0 -p 3306,3307 -S read_write --prepare --tables=32 --table-size=5000000 --threads=64"
    echo ""
    echo "  # 运行测试 (256线程, 持续120秒)"
    echo "  $0 -p 3306,3307 -S read_write,point_select --tables=32 --table-size=5000000 --threads=256 --time=120"
    echo ""
    exit 1
}

# --- 解析参数 ---
PREPARE_DATA="false"
PORTS_STR=""
SCENARIOS_STR=""
# ☆☆☆ 关键改动 1: 创建一个数组来收集所有要传递给 sysbench 的参数 ☆☆☆
SYSBENCH_ARGS=()

# 解析所有参数
while [[ "$#" -gt 0 ]]; do
    case $1 in
        -p) # 脚本的自有参数
            PORTS_STR="$2"
            shift # 跳过值
            ;;
        -S) # 脚本的自有参数
            SCENARIOS_STR="$2"
            shift # 跳过值
            ;;
        --prepare) # 脚本的自有标志
            PREPARE_DATA="true"
            ;;
        --*) # ☆☆☆ 关键改动 2: 捕获所有其他以 '--' 开头的参数 ☆☆☆
            # 这包括 --tables, --table-size, --threads, --time, --rate 等所有 sysbench 参数
            SYSBENCH_ARGS+=("$1")
            ;;
        *)
            echo "错误: 未知或格式错误的参数: $1"
            usage
            ;;
    esac
    shift # 移动到下一个参数
done


if [ -z "$PORTS_STR" ] || [ -z "$SCENARIOS_STR" ]; then
    echo "错误: 必须提供端口列表 (-p) 和场景列表 (-S)。"
    usage
fi

IFS=',' read -r -a PORTS <<< "$PORTS_STR"
IFS=',' read -r -a SCENARIOS <<< "$SCENARIOS_STR"

# --- 函数：为单个实例运行 Sysbench ---
run_sysbench_for_instance() {
    local port=$1
    # ☆☆☆ 关键改动 3: 为了生成唯一的数据库名，我们需要从用户提供的参数中提取 table-size 和 tables ☆☆☆
    # 这是一种健壮的做法，如果用户没提供，我们就用一个默认值。
    local table_size_val=$(echo "${SYSBENCH_ARGS[@]}" | grep -oP '(?<=--table-size=)\d+' || echo "10000")
    local tables_val=$(echo "${SYSBENCH_ARGS[@]}" | grep -oP '(?<=--tables=)\d+' || echo "16")

    local db_name="${DB_NAME_PREFIX}_${table_size_val}x${tables_val}_p${port}"
    local output_dir="${SYSBENCH_OUTPUT_DIR}/port_${port}"
    mkdir -p "$output_dir"

    echo "--- [端口 ${port}] 开始 Sysbench 测试 ---"
    
    # 检查数据库连接并创建库
    mysql -h 127.0.0.1 -P "$port" -u"$MYSQL_USER" -p"$MYSQL_PWD" -e "CREATE DATABASE IF NOT EXISTS \`${db_name}\`;"
    if [ $? -ne 0 ]; then
        echo "错误: [端口 ${port}] 无法连接到 MySQL 或创建数据库。跳过此实例。"
        return 1
    fi

    # 如果指定了 --prepare，则运行准备阶段
    if [ "$PREPARE_DATA" == "true" ]; then
        echo "[端口 ${port}] 正在准备数据 (DB=${db_name})..."
        echo "运行命令: $SYSBENCH_EXECUTABLE ${LUA_SCRIPT_BASE_PATH}/oltp_common.lua --mysql-db=${db_name} ${SYSBENCH_ARGS[@]} prepare"

        "$SYSBENCH_EXECUTABLE" "${LUA_SCRIPT_BASE_PATH}/oltp_common.lua" \
            --mysql-host=127.0.0.1 --mysql-port="$port" \
            --mysql-user="$MYSQL_USER" --mysql-password="$MYSQL_PWD" \
            --mysql-db="$db_name" \
            "${SYSBENCH_ARGS[@]}" \
            prepare
            
        if [ $? -ne 0 ]; then
            echo "错误: [端口 ${port}] 数据准备失败。"
            return 1
        fi
        echo "[端口 ${port}] 数据准备完成。"
    fi

    # 遍历所有场景并运行测试
    for scenario in "${SCENARIOS[@]}"; do
        local lua_script=""
        case "$scenario" in
            read_write)       lua_script="${LUA_SCRIPT_BASE_PATH}/oltp_read_write.lua" ;;
            point_select)     lua_script="${LUA_SCRIPT_BASE_PATH}/oltp_point_select.lua" ;;
            read_only)        lua_script="${LUA_SCRIPT_BASE_PATH}/oltp_read_only.lua" ;;
            write_only)       lua_script="${LUA_SCRIPT_BASE_PATH}/oltp_write_only.lua" ;;
            update_index)     lua_script="${LUA_SCRIPT_BASE_PATH}/oltp_update_index.lua" ;;
            update_non_index) lua_script="${LUA_SCRIPT_BASE_PATH}/oltp_update_non_index.lua" ;;
            *) echo "警告: [端口 ${port}] 不支持的测试场景 '$scenario'。跳过。"; continue ;;
        esac

        if [ ! -f "$lua_script" ]; then
            echo "错误: [端口 ${port}] 找不到场景脚本: $lua_script"; continue
        fi
        
        # 为了日志文件名更具可读性，我们提取一些关键参数
        local threads_val=$(echo "${SYSBENCH_ARGS[@]}" | grep -oP '(?<=--threads=)\d+' || echo "default")
        local time_val=$(echo "${SYSBENCH_ARGS[@]}" | grep -oP '(?<=--time=)\d+' || echo "default")
        local result_log_name="${scenario}_th${threads_val}_time${time_val}_result.log"

        echo "[端口 ${port}] 运行场景: $scenario ..."
        
        # ☆☆☆ 关键改动 4: 在执行 run 命令时，展开参数数组 ☆☆☆
        # "${SYSBENCH_ARGS[@]}" 会将数组中的每个元素作为一个独立的、带引号的参数传递，
        # 完美处理了所有 sysbench 参数。
        "$SYSBENCH_EXECUTABLE" "$lua_script" \
            --mysql-host=127.0.0.1 --mysql-port="$port" \
            --mysql-user="$MYSQL_USER" --mysql-password="$MYSQL_PWD" \
            --mysql-db="$db_name" \
            "${SYSBENCH_ARGS[@]}" \
            run > "${output_dir}/${result_log_name}" 2>&1

        if [ $? -ne 0 ]; then
            echo "错误: [端口 ${port}] 场景 '$scenario' 测试失败。详情请查看: ${output_dir}/${result_log_name}"
        else
            echo "[端口 ${port}] 场景 '$scenario' 测试完成。结果保存在: ${output_dir}/${result_log_name}"
        fi
    done
    return 0
}

# --- 并行运行所有测试 ---
echo "--- 开始并行运行 Sysbench 测试 ---"
echo "端口: ${PORTS_STR}"
echo "场景: ${SCENARIOS_STR}"
echo "Sysbench 附加参数: ${SYSBENCH_ARGS[@]}"
pids=()
for port in "${PORTS[@]}"; do
    # 将 run_sysbench_for_instance 放入后台执行
    run_sysbench_for_instance "$port" &
    pids+=($!)
done

# 等待所有后台任务完成
status=0
for pid in "${pids[@]}"; do
    wait "$pid" || status=1 # 如果任何一个任务失败，将状态设置为1
done

if [ $status -ne 0 ]; then
    echo "警告: 一个或多个 Sysbench 测试未成功完成。"
    exit 1
fi

echo "所有 Sysbench 基准测试已完成。"